import queue
import os
import pickle
from typing import Dict, Any
from stable_baselines3 import PPO

# 自博弈基类 -- 负责存储和管理历史策略
# 使用队列来存储历史策略，最大长度为2048
# 可添加策略、获取指定索引的策略、随机采样策略、获取最新策略、.....
class SelfPlayBase():
    def __init__(self, max_size=2048, temp_storage_dir="./data/temp_self_play"):
        self.history_policy_buffer = queue.Queue(maxsize = max_size)      # 使用队列来存储历史策略，最大长度为2048
        self.temp_storage_dir = temp_storage_dir                          # 临时存储目录
    
        # 确保临时存储目录存在
        os.makedirs(self.temp_storage_dir,exist_ok=True)

    # 添加策略到历史策略缓冲区
    def add_policy(self, policy: PPO) -> None:
        
        if self.history_policy_buffer.full():
            # 如果队列已满，移除最旧的策略存储
            oldest_policy = self.history_policy_buffer.get()
            # 删除对应文件
            os.remove(oldest_policy.policy_file)

        # 生成唯一的策略文件名
        policy_id = f"policy_{len(self.history_policy_buffer)}.pkl"
        policy_file = os.path.join(self.temp_storage_dir,policy_id)
        
        # 保存策略
        policy.save(policy_file)  

        policy.policy_file = policy_file

        # 将文件ID存入队列缓冲区
        self.history_policy_buffer.put(policy)

    # 获取指定索引策略
    def get_index_policy(self, index) -> PPO:
        if 0 <= index < len(self.history_policy_buffer):
            policy = self.history_policy_buffer.queue[index]
            policy_file = policy.policy_file
            # 从文件加载策略
            loaded_policy = PPO.load(policy_file)
            return loaded_policy
        else:
            raise IndexError("Index out of range for history policy buffer.")
        
    # 随机采样策略
    def get_random_policy(self) -> PPO:
        if self.history_policy_buffer:
            import random
            random_policy = random.choice(list(self.history_policy_buffer))
            random_policy_file = random_policy.policy_file
            # 从文件加载策略
            loaded_policy = PPO.load(random_policy_file)
            return loaded_policy
        else:
            raise IndexError("No policies available in the history policy buffer.")

    # 获取最新的策略
    def get_latest_policy(self) -> PPO:
        if not self.history_policy_buffer.empty():
            latest_policy = self.history_policy_buffer.queue[-1]
            return PPO.load(latest_policy.policy_file)
        return None
        
    # 获取所有策略
    def get_all_policies(self) -> list:
         return [policy.policy_file for policy in self.history_policy_buffer.queue]
    
    # 获取策略数量
    def get_policy_count(self) -> int:
        return self.history_policy_buffer.qsize()
    
    # 重置历史策略缓冲区
    def reset(self) -> None:
        # 删除所有临时存储的策略文件
        for policy in self.history_policy_buffer.queue:
            try:
                os.remove(policy.policy_file)
            except FileNotFoundError:
                pass

        self.history_policy_buffer = queue.Queue(maxsize=self.history_policy_buffer.maxsize)  # 重置队列



# 自博弈管理器 -- 管理每个飞机的自博弈过程
class SelfPlayManager():
    def __init__(self, planes_id_list=None,update_steps = 10000):
        
        self.sp_manager = {}        # 用于存储每个飞机的自博弈管理器
        self.sp_steps = {}          # 用于记录每个飞机的自博弈步骤数
        self.cruuent_policy = {}    # 用于存储当前策略
        self.update_steps = update_steps  # 更新自博弈管理器的步数阈值

        # 初始化每个飞机的自博弈管理器
        if planes_id_list is not None:
            for plane_id in planes_id_list:
                self._init_manager(plane_id)

    
    def _init_manager(self, plane_id) -> None:
        if plane_id not in self.sp_manager:
            self.sp_manager[plane_id] = SelfPlayBase(2048, temp_storage_dir=f"./data/temp_self_play/{plane_id}")
            self.sp_steps[plane_id] = 0
            self.cruuent_policy[plane_id] = None

    # 获取指定飞机的动作
    def get_action(self, plane_id, state) -> Any:
        if plane_id in self.sp_manager:

            # 检查是否需要更新自博弈管理器
            if self.sp_steps.get(plane_id, 0) >= self.update_steps:
                # 重新随机采样策略
                new_policy = self.sp_manager[plane_id].get_random_policy()
                self.cruuent_policy[plane_id] = new_policy
                # 重置计步器
                self.sp_steps[plane_id] = 0

            # 如果当前策略不存在，则随机采样一个策略
            if self.cruuent_policy[plane_id] is None:
                # 随机采样动作
                pass

            # 获取当前策略下的动作
            action = self.cruuent_policy[plane_id].predict(state,determine=True)
            self.sp_steps[plane_id] = self.sp_steps.get(plane_id, 0) + 1
            return action
        else:
            raise ValueError(f"Plane ID '{plane_id}' not found in self-play manager.")

    # 更新自博弈管理器
    def update_self_play(self, plane_id:int, policy: PPO) -> None:

        if plane_id not in self.sp_manager:
            self._init_plane_manager(plane_id)
            
        # 添加新策略到自博弈管理器
        self.sp_manager[plane_id].add_policy(policy)
        # 更新当前策略
        self.current_policy[plane_id] = policy
        # 重置计步器
        self.sp_steps[plane_id] = 0

    def add_policy(self, plane_id: int, policy: PPO) -> None:
        """
        添加策略到指定飞机的自博弈管理器
        :param plane_id: 飞机ID
        :param policy: 新策略
        """
        if plane_id in self.sp_manager:
            self.sp_manager[plane_id].add_policy(policy)

    def get_index_policy(self, plane_id: int, index: int) -> PPO:
        """
        获取指定索引的策略
        :param plane_id: 飞机ID
        :param index: 策略索引
        :return: PPO模型实例
        """
        if plane_id in self.sp_manager:
            return self.sp_manager[plane_id].get_index_policy(index)
        return None

    def get_random_policy(self, plane_id: int) -> PPO:
        """
        随机采样策略
        :param plane_id: 飞机ID
        :return: PPO模型实例
        """
        if plane_id in self.sp_manager:
            return self.sp_manager[plane_id].get_random_policy()
        return None

    def get_latest_policy(self, plane_id: int) -> PPO:
        """
        获取最新的策略
        :param plane_id: 飞机ID
        :return: PPO模型实例
        """
        if plane_id in self.sp_manager:
            return self.sp_manager[plane_id].get_latest_policy()
        return None

    def get_all_policies(self, plane_id: int) -> list:
        """
        获取所有策略
        :param plane_id: 飞机ID
        :return: 策略文件路径列表
        """
        if plane_id in self.sp_manager:
            return self.sp_manager[plane_id].get_all_policies()
        return None

    def get_policy_count(self, plane_id: int) -> int:
        """
        获取策略数量
        :param plane_id: 飞机ID
        :return: 策略数量
        """
        if plane_id in self.sp_manager:
            return self.sp_manager[plane_id].get_policy_count()
        return 0

    def clear_policy(self, plane_id: int) -> None:
        """
        清除指定飞机的所有策略
        :param plane_id: 飞机ID
        """
        if plane_id in self.sp_manager:
            self.sp_manager[plane_id].reset()
            self.sp_steps[plane_id] = 0
            self.cruuent_policy[plane_id] = None

